package com.nlp.visualization.core.lda.process;

import com.hankcs.hanlp.corpus.io.IOUtil;
import com.nlp.visualization.core.lda.com.hankcs.lda.Corpus;
import com.nlp.visualization.core.lda.com.hankcs.lda.LdaGibbsSampler;
import com.nlp.visualization.core.lda.com.hankcs.lda.LdaUtil;
import com.nlp.visualization.pojo.lda.LDATheme;
import com.nlp.visualization.service.weibo.main.Theme;
import org.apache.commons.collections.map.HashedMap;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.*;


public class HanLDA {
    LDAModel ldaModel;

    static Logger logger = LoggerFactory.getLogger(HanLDA.class);

    public HanLDA(String modelFilePath) throws IOException {
        ldaModel = new LDAModel(modelFilePath);
    }

    /**
     * 预测文档的主题分布
     *
     * @param documentFilePath 要预测的文档路径
     * @param printResult      是否打印预测结果
     * @return 主题分布数组
     * @throws IOException
     */
    public double[] inference(String documentFilePath, boolean printResult) throws IOException {
        int[] document = Corpus.loadDocument(documentFilePath, ldaModel.getVocabulary());
        //进行文档的主题分布预测
        double[] tp = LdaGibbsSampler.inference(ldaModel.getPhiMatrix(), document);
        if (printResult) {
            for (int i = 0; i < tp.length; i++) {
                System.out.println("topic" + i + "-----" + tp[i]);
            }
        }
        return tp;
    }

    /**
     * 训练模型
     *
     * @param trainFolderPath   训练语料文件所在目录
     * @param trainTopicNum     训练的主题个数
     * @param saveModelFilePath 模型文件保存路径
     * @param printLDAModel     是否打印产生的模型
     * @throws IOException
     * @throws Exception
     */
    public static void train(String trainFolderPath, int trainTopicNum, String saveModelFilePath, boolean printLDAModel) throws IOException {
        //载入语料，预处理
        Corpus corpus = Corpus.load(trainFolderPath);
        //创建LDA采样器
        LdaGibbsSampler ldaGibbsSampler = new LdaGibbsSampler(corpus.getDocument(), corpus.getVocabularySize());
        if (trainTopicNum < 1)
            trainTopicNum = 1; //训练的主题个数至少为1;
        //训练指定个数的主题
        ldaGibbsSampler.gibbs(trainTopicNum);
        //phi矩阵是唯一有用的东西，可以用LdaUtil来展示最终的结果
        double[][] phi = ldaGibbsSampler.getPhi();
        //存储产生的LDA模型
        LDAModel.saveModelFile(saveModelFilePath, phi, corpus.getVocabulary());
        if (printLDAModel) { //展示产生的模型，每个主题展示出现概率最高的10个词汇
            Map<String, Double>[] topicMap = LdaUtil.translate(phi, corpus.getVocabulary(), 20);
            LdaUtil.explain(topicMap);
        }
        //存储LDA展示文件，每个主题展示出现概率最高的20个词汇
        LDAModel.saveLDAShowFile(saveModelFilePath + ".txt", phi, corpus.getVocabulary(), 30);
    }

    /**
     * 预测文档的主题分布
     *
     * @param modelFilePath    模型文件路径
     * @param documentFilePath 要预测的文档路径
     * @param printResult      是否打印预测结果
     * @return 主题分布数组
     * @throws IOException
     */
    public static double[] inference(String modelFilePath, String documentFilePath, boolean printResult) throws IOException {
        LDAModel ldaModel = new LDAModel(modelFilePath);
        int[] document = Corpus.loadDocument(documentFilePath, ldaModel.getVocabulary());
        //进行文档的主题分布预测
        double[] tp = LdaGibbsSampler.inference(ldaModel.getPhiMatrix(), document);
        if (printResult) {
            for (int i = 0; i < tp.length; i++) {
                System.out.println("topic" + i + "-----" + tp[i]);
            }
        }
        return tp;
    }


    /**
     * 预测文档的主题分布,重载
     *
     * @param modelFilePath    模型文件路径
     * @param documentFilePath 要预测的文档路径
     * @return List<LDATheme>
     * @throws IOException
     */
    public static List<LDATheme> inference(String modelFilePath, String documentFilePath) throws IOException {
        List<LDATheme> themeList = new LinkedList<>();
        LDAModel ldaModel = new LDAModel(modelFilePath);
        float others = 1;
        int[] document = Corpus.loadDocument(documentFilePath, ldaModel.getVocabulary());
        //进行文档的主题分布预测
        double[] tp = LdaGibbsSampler.inference(ldaModel.getPhiMatrix(), document);
        //读取主题名称,取出第几天的主题
        String path = modelFilePath.split("\\.")[0] + ".seq";
        File themeNames;
        if (IOUtil.isFileExisted(path))
            themeNames = new File(path);
        else {
            throw new IOException("文件不存在" + path + "");
        }
        List themeNamesList = FileUtils.readLines(themeNames, "utf-8");
        for (int i = 0; i < tp.length; i++) {
            LDATheme theme = new LDATheme();
            theme.setNum(i);
            theme.setProbability((float) tp[i]);
            try {
                theme.setName((String) themeNamesList.get(i));
            } catch (ArrayIndexOutOfBoundsException e) {
                logger.error("主题名称数组出错");
                e.printStackTrace();
                continue;
            }
            themeList.add(theme);

        }


        Collections.sort(themeList, new Comparator<LDATheme>() {
            @Override
            public int compare(LDATheme o1, LDATheme o2) {
                return (int) ((o2.getProbability() - o1.getProbability()) * 1000000000);
            }
        });
        try {
            others -= themeList.get(0).getProbability();
            others -= themeList.get(1).getProbability();
            others -= themeList.get(2).getProbability();
            others -= themeList.get(3).getProbability();
            others -= themeList.get(4).getProbability();
        } catch (ArrayIndexOutOfBoundsException e) {
            logger.error("计算其他的概率时，list下标越界");
            e.printStackTrace();

        }
        LDATheme theme = new LDATheme();
        theme.setName("其他");
        theme.setProbability(others);
        themeList.add(theme);


        Collections.sort(themeList, new Comparator<LDATheme>() {
            @Override
            public int compare(LDATheme o1, LDATheme o2) {
                return (int) ((o2.getProbability() - o1.getProbability()) * 1000000000);
            }
        });

        return themeList;
    }
}
